{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6ebdf92e-deeb-4713-b71b-cae89b64d200",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "70ce09f3-b651-4fd1-8f09-fbdfc75a54d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.read_csv('data/HR.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fe404e64-cdf0-4e8d-bc84-bad8ef3b3da5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>satisfaction_level</th>\n",
       "      <th>last_evaluation</th>\n",
       "      <th>number_project</th>\n",
       "      <th>average_montly_hours</th>\n",
       "      <th>time_spend_company</th>\n",
       "      <th>Work_accident</th>\n",
       "      <th>left</th>\n",
       "      <th>promotion_last_5years</th>\n",
       "      <th>part</th>\n",
       "      <th>salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.38</td>\n",
       "      <td>0.53</td>\n",
       "      <td>2</td>\n",
       "      <td>157</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.80</td>\n",
       "      <td>0.86</td>\n",
       "      <td>5</td>\n",
       "      <td>262</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>medium</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.11</td>\n",
       "      <td>0.88</td>\n",
       "      <td>7</td>\n",
       "      <td>272</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>medium</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.72</td>\n",
       "      <td>0.87</td>\n",
       "      <td>5</td>\n",
       "      <td>223</td>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.52</td>\n",
       "      <td>2</td>\n",
       "      <td>159</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   satisfaction_level  last_evaluation  number_project  average_montly_hours  \\\n",
       "0                0.38             0.53               2                   157   \n",
       "1                0.80             0.86               5                   262   \n",
       "2                0.11             0.88               7                   272   \n",
       "3                0.72             0.87               5                   223   \n",
       "4                0.37             0.52               2                   159   \n",
       "\n",
       "   time_spend_company  Work_accident  left  promotion_last_5years   part  \\\n",
       "0                   3              0     1                      0  sales   \n",
       "1                   6              0     1                      0  sales   \n",
       "2                   4              0     1                      0  sales   \n",
       "3                   5              0     1                      0  sales   \n",
       "4                   3              0     1                      0  sales   \n",
       "\n",
       "   salary  \n",
       "0     low  \n",
       "1  medium  \n",
       "2  medium  \n",
       "3     low  \n",
       "4     low  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "66da4127-30b5-4350-b3e7-92d106e9c6a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 14999 entries, 0 to 14998\n",
      "Data columns (total 10 columns):\n",
      " #   Column                 Non-Null Count  Dtype  \n",
      "---  ------                 --------------  -----  \n",
      " 0   satisfaction_level     14999 non-null  float64\n",
      " 1   last_evaluation        14999 non-null  float64\n",
      " 2   number_project         14999 non-null  int64  \n",
      " 3   average_montly_hours   14999 non-null  int64  \n",
      " 4   time_spend_company     14999 non-null  int64  \n",
      " 5   Work_accident          14999 non-null  int64  \n",
      " 6   left                   14999 non-null  int64  \n",
      " 7   promotion_last_5years  14999 non-null  int64  \n",
      " 8   part                   14999 non-null  object \n",
      " 9   salary                 14999 non-null  object \n",
      "dtypes: float64(2), int64(6), object(2)\n",
      "memory usage: 1.1+ MB\n"
     ]
    }
   ],
   "source": [
    "# 查看表格数据类型\n",
    "data.info() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9bed4b43-ef7f-4eaa-a6a4-793204aafa1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['sales', 'accounting', 'hr', 'technical', 'support', 'management',\n",
       "       'IT', 'product_mng', 'marketing', 'RandD'], dtype=object)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.part.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a6aad9e3-4582-40b4-8adf-9cecbb84f299",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['low', 'medium', 'high'], dtype=object)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.salary.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8e1759d8-da73-4300-818f-8f3e8d129a59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "salary  part       \n",
       "high    IT               83\n",
       "        RandD            51\n",
       "        accounting       74\n",
       "        hr               45\n",
       "        management      225\n",
       "        marketing        80\n",
       "        product_mng      68\n",
       "        sales           269\n",
       "        support         141\n",
       "        technical       201\n",
       "low     IT              609\n",
       "        RandD           364\n",
       "        accounting      358\n",
       "        hr              335\n",
       "        management      180\n",
       "        marketing       402\n",
       "        product_mng     451\n",
       "        sales          2099\n",
       "        support        1146\n",
       "        technical      1372\n",
       "medium  IT              535\n",
       "        RandD           372\n",
       "        accounting      335\n",
       "        hr              359\n",
       "        management      225\n",
       "        marketing       376\n",
       "        product_mng     383\n",
       "        sales          1772\n",
       "        support         942\n",
       "        technical      1147\n",
       "dtype: int64"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.groupby(['salary','part']).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c11af8d3-0269-4e3c-952f-a1338c2b7885",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>high</th>\n",
       "      <th>low</th>\n",
       "      <th>medium</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14994</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14995</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14996</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14997</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14998</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>14999 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       high  low  medium\n",
       "0         0    1       0\n",
       "1         0    0       1\n",
       "2         0    0       1\n",
       "3         0    1       0\n",
       "4         0    1       0\n",
       "...     ...  ...     ...\n",
       "14994     0    1       0\n",
       "14995     0    1       0\n",
       "14996     0    1       0\n",
       "14997     0    1       0\n",
       "14998     0    1       0\n",
       "\n",
       "[14999 rows x 3 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.get_dummies(data.salary,dtype = int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "35e40691-7b7a-40b9-8c1c-337d0d18252d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将salary展开添加到data的最后\n",
    "data = data.join(pd.get_dummies(data.salary,dtype=int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2cfe27f2-70c5-42b2-890b-e816c1e73b01",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.join(pd.get_dummies(data.part,dtype=int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ecc65f84-312c-4248-9a08-79bc2dedc833",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>satisfaction_level</th>\n",
       "      <th>last_evaluation</th>\n",
       "      <th>number_project</th>\n",
       "      <th>average_montly_hours</th>\n",
       "      <th>time_spend_company</th>\n",
       "      <th>Work_accident</th>\n",
       "      <th>left</th>\n",
       "      <th>promotion_last_5years</th>\n",
       "      <th>part</th>\n",
       "      <th>salary</th>\n",
       "      <th>...</th>\n",
       "      <th>IT</th>\n",
       "      <th>RandD</th>\n",
       "      <th>accounting</th>\n",
       "      <th>hr</th>\n",
       "      <th>management</th>\n",
       "      <th>marketing</th>\n",
       "      <th>product_mng</th>\n",
       "      <th>sales</th>\n",
       "      <th>support</th>\n",
       "      <th>technical</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.38</td>\n",
       "      <td>0.53</td>\n",
       "      <td>2</td>\n",
       "      <td>157</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.80</td>\n",
       "      <td>0.86</td>\n",
       "      <td>5</td>\n",
       "      <td>262</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>medium</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.11</td>\n",
       "      <td>0.88</td>\n",
       "      <td>7</td>\n",
       "      <td>272</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>medium</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.72</td>\n",
       "      <td>0.87</td>\n",
       "      <td>5</td>\n",
       "      <td>223</td>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.52</td>\n",
       "      <td>2</td>\n",
       "      <td>159</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>sales</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14994</th>\n",
       "      <td>0.40</td>\n",
       "      <td>0.57</td>\n",
       "      <td>2</td>\n",
       "      <td>151</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>support</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14995</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.48</td>\n",
       "      <td>2</td>\n",
       "      <td>160</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>support</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14996</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.53</td>\n",
       "      <td>2</td>\n",
       "      <td>143</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>support</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14997</th>\n",
       "      <td>0.11</td>\n",
       "      <td>0.96</td>\n",
       "      <td>6</td>\n",
       "      <td>280</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>support</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14998</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.52</td>\n",
       "      <td>2</td>\n",
       "      <td>158</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>support</td>\n",
       "      <td>low</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>14999 rows × 23 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       satisfaction_level  last_evaluation  number_project  \\\n",
       "0                    0.38             0.53               2   \n",
       "1                    0.80             0.86               5   \n",
       "2                    0.11             0.88               7   \n",
       "3                    0.72             0.87               5   \n",
       "4                    0.37             0.52               2   \n",
       "...                   ...              ...             ...   \n",
       "14994                0.40             0.57               2   \n",
       "14995                0.37             0.48               2   \n",
       "14996                0.37             0.53               2   \n",
       "14997                0.11             0.96               6   \n",
       "14998                0.37             0.52               2   \n",
       "\n",
       "       average_montly_hours  time_spend_company  Work_accident  left  \\\n",
       "0                       157                   3              0     1   \n",
       "1                       262                   6              0     1   \n",
       "2                       272                   4              0     1   \n",
       "3                       223                   5              0     1   \n",
       "4                       159                   3              0     1   \n",
       "...                     ...                 ...            ...   ...   \n",
       "14994                   151                   3              0     1   \n",
       "14995                   160                   3              0     1   \n",
       "14996                   143                   3              0     1   \n",
       "14997                   280                   4              0     1   \n",
       "14998                   158                   3              0     1   \n",
       "\n",
       "       promotion_last_5years     part  salary  ...  IT  RandD  accounting  hr  \\\n",
       "0                          0    sales     low  ...   0      0           0   0   \n",
       "1                          0    sales  medium  ...   0      0           0   0   \n",
       "2                          0    sales  medium  ...   0      0           0   0   \n",
       "3                          0    sales     low  ...   0      0           0   0   \n",
       "4                          0    sales     low  ...   0      0           0   0   \n",
       "...                      ...      ...     ...  ...  ..    ...         ...  ..   \n",
       "14994                      0  support     low  ...   0      0           0   0   \n",
       "14995                      0  support     low  ...   0      0           0   0   \n",
       "14996                      0  support     low  ...   0      0           0   0   \n",
       "14997                      0  support     low  ...   0      0           0   0   \n",
       "14998                      0  support     low  ...   0      0           0   0   \n",
       "\n",
       "       management  marketing  product_mng  sales  support  technical  \n",
       "0               0          0            0      1        0          0  \n",
       "1               0          0            0      1        0          0  \n",
       "2               0          0            0      1        0          0  \n",
       "3               0          0            0      1        0          0  \n",
       "4               0          0            0      1        0          0  \n",
       "...           ...        ...          ...    ...      ...        ...  \n",
       "14994           0          0            0      0        1          0  \n",
       "14995           0          0            0      0        1          0  \n",
       "14996           0          0            0      0        1          0  \n",
       "14997           0          0            0      0        1          0  \n",
       "14998           0          0            0      0        1          0  \n",
       "\n",
       "[14999 rows x 23 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "01352501-4015-46df-9de5-684d59f41826",
   "metadata": {},
   "outputs": [],
   "source": [
    "del data['salary']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f48f28f9-5fa0-4b57-adde-b732fcc2f748",
   "metadata": {},
   "outputs": [],
   "source": [
    "del data['part']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "60f99e52-ccfc-4e04-b906-a515feb6e2a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>satisfaction_level</th>\n",
       "      <th>last_evaluation</th>\n",
       "      <th>number_project</th>\n",
       "      <th>average_montly_hours</th>\n",
       "      <th>time_spend_company</th>\n",
       "      <th>Work_accident</th>\n",
       "      <th>left</th>\n",
       "      <th>promotion_last_5years</th>\n",
       "      <th>high</th>\n",
       "      <th>low</th>\n",
       "      <th>...</th>\n",
       "      <th>IT</th>\n",
       "      <th>RandD</th>\n",
       "      <th>accounting</th>\n",
       "      <th>hr</th>\n",
       "      <th>management</th>\n",
       "      <th>marketing</th>\n",
       "      <th>product_mng</th>\n",
       "      <th>sales</th>\n",
       "      <th>support</th>\n",
       "      <th>technical</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.38</td>\n",
       "      <td>0.53</td>\n",
       "      <td>2</td>\n",
       "      <td>157</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.80</td>\n",
       "      <td>0.86</td>\n",
       "      <td>5</td>\n",
       "      <td>262</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.11</td>\n",
       "      <td>0.88</td>\n",
       "      <td>7</td>\n",
       "      <td>272</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.72</td>\n",
       "      <td>0.87</td>\n",
       "      <td>5</td>\n",
       "      <td>223</td>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.52</td>\n",
       "      <td>2</td>\n",
       "      <td>159</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14994</th>\n",
       "      <td>0.40</td>\n",
       "      <td>0.57</td>\n",
       "      <td>2</td>\n",
       "      <td>151</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14995</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.48</td>\n",
       "      <td>2</td>\n",
       "      <td>160</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14996</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.53</td>\n",
       "      <td>2</td>\n",
       "      <td>143</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14997</th>\n",
       "      <td>0.11</td>\n",
       "      <td>0.96</td>\n",
       "      <td>6</td>\n",
       "      <td>280</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14998</th>\n",
       "      <td>0.37</td>\n",
       "      <td>0.52</td>\n",
       "      <td>2</td>\n",
       "      <td>158</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>14999 rows × 21 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       satisfaction_level  last_evaluation  number_project  \\\n",
       "0                    0.38             0.53               2   \n",
       "1                    0.80             0.86               5   \n",
       "2                    0.11             0.88               7   \n",
       "3                    0.72             0.87               5   \n",
       "4                    0.37             0.52               2   \n",
       "...                   ...              ...             ...   \n",
       "14994                0.40             0.57               2   \n",
       "14995                0.37             0.48               2   \n",
       "14996                0.37             0.53               2   \n",
       "14997                0.11             0.96               6   \n",
       "14998                0.37             0.52               2   \n",
       "\n",
       "       average_montly_hours  time_spend_company  Work_accident  left  \\\n",
       "0                       157                   3              0     1   \n",
       "1                       262                   6              0     1   \n",
       "2                       272                   4              0     1   \n",
       "3                       223                   5              0     1   \n",
       "4                       159                   3              0     1   \n",
       "...                     ...                 ...            ...   ...   \n",
       "14994                   151                   3              0     1   \n",
       "14995                   160                   3              0     1   \n",
       "14996                   143                   3              0     1   \n",
       "14997                   280                   4              0     1   \n",
       "14998                   158                   3              0     1   \n",
       "\n",
       "       promotion_last_5years  high  low  ...  IT  RandD  accounting  hr  \\\n",
       "0                          0     0    1  ...   0      0           0   0   \n",
       "1                          0     0    0  ...   0      0           0   0   \n",
       "2                          0     0    0  ...   0      0           0   0   \n",
       "3                          0     0    1  ...   0      0           0   0   \n",
       "4                          0     0    1  ...   0      0           0   0   \n",
       "...                      ...   ...  ...  ...  ..    ...         ...  ..   \n",
       "14994                      0     0    1  ...   0      0           0   0   \n",
       "14995                      0     0    1  ...   0      0           0   0   \n",
       "14996                      0     0    1  ...   0      0           0   0   \n",
       "14997                      0     0    1  ...   0      0           0   0   \n",
       "14998                      0     0    1  ...   0      0           0   0   \n",
       "\n",
       "       management  marketing  product_mng  sales  support  technical  \n",
       "0               0          0            0      1        0          0  \n",
       "1               0          0            0      1        0          0  \n",
       "2               0          0            0      1        0          0  \n",
       "3               0          0            0      1        0          0  \n",
       "4               0          0            0      1        0          0  \n",
       "...           ...        ...          ...    ...      ...        ...  \n",
       "14994           0          0            0      0        1          0  \n",
       "14995           0          0            0      0        1          0  \n",
       "14996           0          0            0      0        1          0  \n",
       "14997           0          0            0      0        1          0  \n",
       "14998           0          0            0      0        1          0  \n",
       "\n",
       "[14999 rows x 21 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "22783152-3d65-41b4-8813-8a906a172acb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "left\n",
       "0    11428\n",
       "1     3571\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.left.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1924e681-4cc4-4bf4-944f-8bb9bb791549",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = torch.from_numpy(data.left.values.reshape(-1,1)).type(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fdb6797a-0b3b-4328-85f3-ad9526331a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_data = data[[v for v in data.columns if v != 'left']].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a97464cd-8faa-4137-ad46-3516faf131d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.from_numpy(X_data).type(torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "44646a58-c85e-4ef4-9bec-1670cb4d8613",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([14999, 20])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2b62ab7-7d0e-4a5b-b998-229140ad7b68",
   "metadata": {},
   "source": [
    "# 创建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "178f254e-942e-45a4-8c8f-2515dc39ba5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32dcf122-0191-437a-be98-a599780e1d84",
   "metadata": {},
   "source": [
    "自定义模型:\n",
    "nn.Module 都继承这个类\n",
    "__init__: 初始化所有的层\n",
    "forward : 定义模型的运算过程（前向传播的过程）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "3fd2d7dd-739b-413a-8d74-7102dc274e05",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(20,64)\n",
    "        self.linear2 = nn.Linear(64,64)\n",
    "        self.linear3 = nn.Linear(64,1)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "    def forward(self,input):\n",
    "        x = self.linear1(input)\n",
    "        x = self.relu(x)\n",
    "        x = self.linear2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.linear3(x)\n",
    "        x = self.sigmoid(x)\n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "50a5427a-83a1-4392-b838-759049e0140a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "58455ce3-f8a7-4219-8059-5ad0fb4ab91c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (linear1): Linear(in_features=20, out_features=64, bias=True)\n",
       "  (linear2): Linear(in_features=64, out_features=64, bias=True)\n",
       "  (linear3): Linear(in_features=64, out_features=1, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (sigmoid): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f41e82b-1321-40f8-8775-46ddb2d80e8b",
   "metadata": {},
   "source": [
    "# 用functional改写"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "8db0b8b9-ffd0-41f8-8c9f-16415ef916c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "71d9b2d8-6b20-48a7-ad29-d567f1b594b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(20,64)\n",
    "        self.linear2 = nn.Linear(64,64)\n",
    "        self.linear3 = nn.Linear(64,1)\n",
    "    def forward(self,input):\n",
    "        x = F.relu(self.linear1(input)) \n",
    "        x = F.relu(self.linear2(x))\n",
    "        x = F.sigmoid(self.linear3(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "90c6f033-52f7-4429-b50b-4175525ab611",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "bbfb9cb2-a1a3-44b8-89fd-81e8bba1072a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (linear1): Linear(in_features=20, out_features=64, bias=True)\n",
       "  (linear2): Linear(in_features=64, out_features=64, bias=True)\n",
       "  (linear3): Linear(in_features=64, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "754292da-2fa7-4633-b9e6-c0c0acdfd8ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.0001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "0fd33c1c-492f-4f0a-8f5c-94b60d55ee2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    model = Model()\n",
    "    opt = torch.optim.Adam(model.parameters(),lr=lr)\n",
    "    return model,opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "a5eaa555-1618-4931-aa90-90e7f2ee6c55",
   "metadata": {},
   "outputs": [],
   "source": [
    "model,optim = get_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60d3484-080c-4905-866a-93c09c93eabb",
   "metadata": {},
   "source": [
    "## 定义损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "86a1ff38-8f98-47b9-b352-98f18fa4116f",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5cb05d78-cae0-4b30-8084-afe86d4acd43",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "0d13aadc-900f-4c2b-85ab-03d6a8ee6b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "no_of_batches = len(data) // batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "30b20d02-d8bd-4b08-8852-855c568f5562",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "fae98188-c24e-49ff-b6f1-2507db0aff35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:  0 loss:  0.51776522397995\n",
      "epoch:  1 loss:  0.5146937966346741\n",
      "epoch:  2 loss:  0.5129804611206055\n",
      "epoch:  3 loss:  0.510653555393219\n",
      "epoch:  4 loss:  0.5098510980606079\n",
      "epoch:  5 loss:  0.5077739953994751\n",
      "epoch:  6 loss:  0.5072493553161621\n",
      "epoch:  7 loss:  0.505224883556366\n",
      "epoch:  8 loss:  0.5045336484909058\n",
      "epoch:  9 loss:  0.502454936504364\n",
      "epoch:  10 loss:  0.5016345381736755\n",
      "epoch:  11 loss:  0.4996410012245178\n",
      "epoch:  12 loss:  0.49895599484443665\n",
      "epoch:  13 loss:  0.4971199631690979\n",
      "epoch:  14 loss:  0.4962589740753174\n",
      "epoch:  15 loss:  0.4943163990974426\n",
      "epoch:  16 loss:  0.4960300624370575\n",
      "epoch:  17 loss:  0.49825507402420044\n",
      "epoch:  18 loss:  0.4926063120365143\n",
      "epoch:  19 loss:  0.4895774722099304\n",
      "epoch:  20 loss:  0.48780152201652527\n",
      "epoch:  21 loss:  0.4868218004703522\n",
      "epoch:  22 loss:  0.4851723313331604\n",
      "epoch:  23 loss:  0.48493558168411255\n",
      "epoch:  24 loss:  0.48344647884368896\n",
      "epoch:  25 loss:  0.4828924238681793\n",
      "epoch:  26 loss:  0.4811307191848755\n",
      "epoch:  27 loss:  0.48046037554740906\n",
      "epoch:  28 loss:  0.47844675183296204\n",
      "epoch:  29 loss:  0.47771957516670227\n",
      "epoch:  30 loss:  0.4768574833869934\n",
      "epoch:  31 loss:  0.47562727332115173\n",
      "epoch:  32 loss:  0.4746188819408417\n",
      "epoch:  33 loss:  0.4734484851360321\n",
      "epoch:  34 loss:  0.47304996848106384\n",
      "epoch:  35 loss:  0.4719056785106659\n",
      "epoch:  36 loss:  0.4702702760696411\n",
      "epoch:  37 loss:  0.47096750140190125\n",
      "epoch:  38 loss:  0.4680165648460388\n",
      "epoch:  39 loss:  0.4674440622329712\n",
      "epoch:  40 loss:  0.4659419357776642\n",
      "epoch:  41 loss:  0.465438574552536\n",
      "epoch:  42 loss:  0.4635453522205353\n",
      "epoch:  43 loss:  0.463623583316803\n",
      "epoch:  44 loss:  0.46197396516799927\n",
      "epoch:  45 loss:  0.4612341523170471\n",
      "epoch:  46 loss:  0.4596388339996338\n",
      "epoch:  47 loss:  0.46025899052619934\n",
      "epoch:  48 loss:  0.4574282169342041\n",
      "epoch:  49 loss:  0.4572705328464508\n",
      "epoch:  50 loss:  0.4552759528160095\n",
      "epoch:  51 loss:  0.45473214983940125\n",
      "epoch:  52 loss:  0.4525116980075836\n",
      "epoch:  53 loss:  0.45380744338035583\n",
      "epoch:  54 loss:  0.45164549350738525\n",
      "epoch:  55 loss:  0.45019784569740295\n",
      "epoch:  56 loss:  0.44879019260406494\n",
      "epoch:  57 loss:  0.4477252960205078\n",
      "epoch:  58 loss:  0.44375506043434143\n",
      "epoch:  59 loss:  0.44407010078430176\n",
      "epoch:  60 loss:  0.44046536087989807\n",
      "epoch:  61 loss:  0.4410971701145172\n",
      "epoch:  62 loss:  0.43808889389038086\n",
      "epoch:  63 loss:  0.43615037202835083\n",
      "epoch:  64 loss:  0.4356433153152466\n",
      "epoch:  65 loss:  0.436908483505249\n",
      "epoch:  66 loss:  0.4332744777202606\n",
      "epoch:  67 loss:  0.4302773177623749\n",
      "epoch:  68 loss:  0.4275832772254944\n",
      "epoch:  69 loss:  0.42758458852767944\n",
      "epoch:  70 loss:  0.42468318343162537\n",
      "epoch:  71 loss:  0.4233410060405731\n",
      "epoch:  72 loss:  0.42112982273101807\n",
      "epoch:  73 loss:  0.421098917722702\n",
      "epoch:  74 loss:  0.4165334403514862\n",
      "epoch:  75 loss:  0.4161214530467987\n",
      "epoch:  76 loss:  0.4122427701950073\n",
      "epoch:  77 loss:  0.41156649589538574\n",
      "epoch:  78 loss:  0.40990903973579407\n",
      "epoch:  79 loss:  0.4076181948184967\n",
      "epoch:  80 loss:  0.40699151158332825\n",
      "epoch:  81 loss:  0.40324968099594116\n",
      "epoch:  82 loss:  0.4107436239719391\n",
      "epoch:  83 loss:  0.4066057801246643\n",
      "epoch:  84 loss:  0.403799831867218\n",
      "epoch:  85 loss:  0.3999207019805908\n",
      "epoch:  86 loss:  0.3992849588394165\n",
      "epoch:  87 loss:  0.39742565155029297\n",
      "epoch:  88 loss:  0.39561977982521057\n",
      "epoch:  89 loss:  0.3940674066543579\n",
      "epoch:  90 loss:  0.39231789112091064\n",
      "epoch:  91 loss:  0.39111220836639404\n",
      "epoch:  92 loss:  0.3889680504798889\n",
      "epoch:  93 loss:  0.38797447085380554\n",
      "epoch:  94 loss:  0.38629472255706787\n",
      "epoch:  95 loss:  0.3851744830608368\n",
      "epoch:  96 loss:  0.38361525535583496\n",
      "epoch:  97 loss:  0.3823276460170746\n",
      "epoch:  98 loss:  0.38114339113235474\n",
      "epoch:  99 loss:  0.37996068596839905\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(no_of_batches):\n",
    "        start = i * batch\n",
    "        end = start + batch\n",
    "        x = X[start:end]\n",
    "        y = Y[start:end]\n",
    "        y_gred = model(x)\n",
    "        loss = loss_fn(y_gred,y)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "    with torch.no_grad():\n",
    "        print('epoch: ',epoch,'loss: ',loss_fn(model(X),Y).data.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911e3f30-7776-42f1-84b6-92a8ed794188",
   "metadata": {},
   "source": [
    "# 使用dataset重构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4c1fc12e-fb22-4a35-9d9f-01f83214c73c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "502984f6-de9f-4b95-9bbb-c0b0168a789a",
   "metadata": {},
   "outputs": [],
   "source": [
    "HRdataset = TensorDataset(X,Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "d5953a16-e071-440d-b776-6da2be7e2b96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14999"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(HRdataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "a5314135-1290-4229-8a80-cf1e06577693",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([  0.3800,   0.5300,   2.0000, 157.0000,   3.0000,   0.0000,   0.0000,\n",
       "           0.0000,   1.0000,   0.0000,   0.0000,   0.0000,   0.0000,   0.0000,\n",
       "           0.0000,   0.0000,   0.0000,   1.0000,   0.0000,   0.0000]),\n",
       " tensor([1.]))"
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "HRdataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "f26cd7a1-4ab9-4558-86c6-df5f9c30d978",
   "metadata": {},
   "outputs": [],
   "source": [
    "model,optim = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "77ade111-f5e7-4c8a-96cb-5b6337a0e29a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:  0 loss:  0.5132989883422852\n",
      "epoch:  1 loss:  0.5118004679679871\n",
      "epoch:  2 loss:  0.5107380747795105\n",
      "epoch:  3 loss:  0.5092574954032898\n",
      "epoch:  4 loss:  0.5079694986343384\n",
      "epoch:  5 loss:  0.5067866444587708\n",
      "epoch:  6 loss:  0.5055269598960876\n",
      "epoch:  7 loss:  0.5041393041610718\n",
      "epoch:  8 loss:  0.5025852918624878\n",
      "epoch:  9 loss:  0.5012056827545166\n",
      "epoch:  10 loss:  0.499681681394577\n",
      "epoch:  11 loss:  0.49908342957496643\n",
      "epoch:  12 loss:  0.4970693588256836\n",
      "epoch:  13 loss:  0.4962054193019867\n",
      "epoch:  14 loss:  0.4945310354232788\n",
      "epoch:  15 loss:  0.4916222095489502\n",
      "epoch:  16 loss:  0.49351227283477783\n",
      "epoch:  17 loss:  0.4864089787006378\n",
      "epoch:  18 loss:  0.48382705450057983\n",
      "epoch:  19 loss:  0.4839892089366913\n",
      "epoch:  20 loss:  0.47996070981025696\n",
      "epoch:  21 loss:  0.4778911769390106\n",
      "epoch:  22 loss:  0.47638776898384094\n",
      "epoch:  23 loss:  0.47506338357925415\n",
      "epoch:  24 loss:  0.47335371375083923\n",
      "epoch:  25 loss:  0.4717104732990265\n",
      "epoch:  26 loss:  0.46943074464797974\n",
      "epoch:  27 loss:  0.4678933918476105\n",
      "epoch:  28 loss:  0.4642283320426941\n",
      "epoch:  29 loss:  0.46240517497062683\n",
      "epoch:  30 loss:  0.45916813611984253\n",
      "epoch:  31 loss:  0.45828112959861755\n",
      "epoch:  32 loss:  0.4553838074207306\n",
      "epoch:  33 loss:  0.45479172468185425\n",
      "epoch:  34 loss:  0.45211073756217957\n",
      "epoch:  35 loss:  0.45042702555656433\n",
      "epoch:  36 loss:  0.4699118435382843\n",
      "epoch:  37 loss:  0.45103907585144043\n",
      "epoch:  38 loss:  0.44896218180656433\n",
      "epoch:  39 loss:  0.4464438259601593\n",
      "epoch:  40 loss:  0.4433383345603943\n",
      "epoch:  41 loss:  0.4430307447910309\n",
      "epoch:  42 loss:  0.4403670132160187\n",
      "epoch:  43 loss:  0.43883588910102844\n",
      "epoch:  44 loss:  0.43658044934272766\n",
      "epoch:  45 loss:  0.43459901213645935\n",
      "epoch:  46 loss:  0.433459997177124\n",
      "epoch:  47 loss:  0.431903213262558\n",
      "epoch:  48 loss:  0.45978328585624695\n",
      "epoch:  49 loss:  0.45060768723487854\n",
      "epoch:  50 loss:  0.43950843811035156\n",
      "epoch:  51 loss:  0.43800970911979675\n",
      "epoch:  52 loss:  0.4351452887058258\n",
      "epoch:  53 loss:  0.4336109757423401\n",
      "epoch:  54 loss:  0.43191206455230713\n",
      "epoch:  55 loss:  0.4296228289604187\n",
      "epoch:  56 loss:  0.42860081791877747\n",
      "epoch:  57 loss:  0.42718690633773804\n",
      "epoch:  58 loss:  0.4253710210323334\n",
      "epoch:  59 loss:  0.42367109656333923\n",
      "epoch:  60 loss:  0.42225411534309387\n",
      "epoch:  61 loss:  0.4202166199684143\n",
      "epoch:  62 loss:  0.41998690366744995\n",
      "epoch:  63 loss:  0.41770419478416443\n",
      "epoch:  64 loss:  0.4150426983833313\n",
      "epoch:  65 loss:  0.4119613468647003\n",
      "epoch:  66 loss:  0.40938639640808105\n",
      "epoch:  67 loss:  0.40690305829048157\n",
      "epoch:  68 loss:  0.40491563081741333\n",
      "epoch:  69 loss:  0.4028986692428589\n",
      "epoch:  70 loss:  0.40111473202705383\n",
      "epoch:  71 loss:  0.39911046624183655\n",
      "epoch:  72 loss:  0.39775151014328003\n",
      "epoch:  73 loss:  0.3960410952568054\n",
      "epoch:  74 loss:  0.39429011940956116\n",
      "epoch:  75 loss:  0.3925943970680237\n",
      "epoch:  76 loss:  0.3910086452960968\n",
      "epoch:  77 loss:  0.38929834961891174\n",
      "epoch:  78 loss:  0.3877711892127991\n",
      "epoch:  79 loss:  0.38602373003959656\n",
      "epoch:  80 loss:  0.3844839632511139\n",
      "epoch:  81 loss:  0.38287031650543213\n",
      "epoch:  82 loss:  0.38136082887649536\n",
      "epoch:  83 loss:  0.38000449538230896\n",
      "epoch:  84 loss:  0.37890923023223877\n",
      "epoch:  85 loss:  0.3776209354400635\n",
      "epoch:  86 loss:  0.37658658623695374\n",
      "epoch:  87 loss:  0.4028858542442322\n",
      "epoch:  88 loss:  0.4019959270954132\n",
      "epoch:  89 loss:  0.3878229260444641\n",
      "epoch:  90 loss:  0.38172629475593567\n",
      "epoch:  91 loss:  0.3779734969139099\n",
      "epoch:  92 loss:  0.37455108761787415\n",
      "epoch:  93 loss:  0.3723478317260742\n",
      "epoch:  94 loss:  0.37141382694244385\n",
      "epoch:  95 loss:  0.3700305223464966\n",
      "epoch:  96 loss:  0.36925455927848816\n",
      "epoch:  97 loss:  0.36733385920524597\n",
      "epoch:  98 loss:  0.36610037088394165\n",
      "epoch:  99 loss:  0.36406949162483215\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(no_of_batches):\n",
    "        # start = i * batch\n",
    "        # end = start + batch\n",
    "        # x = X[start:end]\n",
    "        # y = Y[start:end]\n",
    "        # 使用这个替换调上面的啰嗦的逻辑\n",
    "        x,y = HRdataset[i * batch:i * batch + batch]\n",
    "        y_gred = model(x)\n",
    "        loss = loss_fn(y_gred,y)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "    with torch.no_grad():\n",
    "        print('epoch: ',epoch,'loss: ',loss_fn(model(X),Y).data.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cc21945-00de-4960-ace8-cfe81609edf3",
   "metadata": {},
   "source": [
    "# 使用dataload 重构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "7f0c3db0-3461-484f-ad80-bd3f176d07bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "c82b90be-a1e0-4b93-84d4-123fd82522a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "HR_ds = TensorDataset(X,Y)\n",
    "HR_dl = DataLoader(HR_ds,batch_size=batch,shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "0dc888bc-42dd-41d5-b977-bfa09ee1757d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:  0 loss:  0.35045361518859863\n",
      "epoch:  1 loss:  0.34200868010520935\n",
      "epoch:  2 loss:  0.3336351215839386\n",
      "epoch:  3 loss:  0.32581740617752075\n",
      "epoch:  4 loss:  0.31905075907707214\n",
      "epoch:  5 loss:  0.3104759156703949\n",
      "epoch:  6 loss:  0.3051707446575165\n",
      "epoch:  7 loss:  0.29744598269462585\n",
      "epoch:  8 loss:  0.2912115752696991\n",
      "epoch:  9 loss:  0.2875969409942627\n",
      "epoch:  10 loss:  0.28045111894607544\n",
      "epoch:  11 loss:  0.2758060395717621\n",
      "epoch:  12 loss:  0.27096107602119446\n",
      "epoch:  13 loss:  0.2659701108932495\n",
      "epoch:  14 loss:  0.2630513608455658\n",
      "epoch:  15 loss:  0.25773558020591736\n",
      "epoch:  16 loss:  0.2542820870876312\n",
      "epoch:  17 loss:  0.25057923793792725\n",
      "epoch:  18 loss:  0.2474661022424698\n",
      "epoch:  19 loss:  0.24427178502082825\n",
      "epoch:  20 loss:  0.24061037600040436\n",
      "epoch:  21 loss:  0.23835790157318115\n",
      "epoch:  22 loss:  0.23614980280399323\n",
      "epoch:  23 loss:  0.2332720160484314\n",
      "epoch:  24 loss:  0.23171064257621765\n",
      "epoch:  25 loss:  0.22946877777576447\n",
      "epoch:  26 loss:  0.22804628312587738\n",
      "epoch:  27 loss:  0.22514979541301727\n",
      "epoch:  28 loss:  0.22449235618114471\n",
      "epoch:  29 loss:  0.22213266789913177\n",
      "epoch:  30 loss:  0.2204238474369049\n",
      "epoch:  31 loss:  0.21950751543045044\n",
      "epoch:  32 loss:  0.21860475838184357\n",
      "epoch:  33 loss:  0.2171705663204193\n",
      "epoch:  34 loss:  0.21512965857982635\n",
      "epoch:  35 loss:  0.21717676520347595\n",
      "epoch:  36 loss:  0.2142612636089325\n",
      "epoch:  37 loss:  0.21234086155891418\n",
      "epoch:  38 loss:  0.2172578126192093\n",
      "epoch:  39 loss:  0.21597172319889069\n",
      "epoch:  40 loss:  0.21016088128089905\n",
      "epoch:  41 loss:  0.20903098583221436\n",
      "epoch:  42 loss:  0.2089749127626419\n",
      "epoch:  43 loss:  0.20942121744155884\n",
      "epoch:  44 loss:  0.20695684850215912\n",
      "epoch:  45 loss:  0.20690219104290009\n",
      "epoch:  46 loss:  0.20652320981025696\n",
      "epoch:  47 loss:  0.2065633088350296\n",
      "epoch:  48 loss:  0.2077174335718155\n",
      "epoch:  49 loss:  0.2050962597131729\n",
      "epoch:  50 loss:  0.205500528216362\n",
      "epoch:  51 loss:  0.20719417929649353\n",
      "epoch:  52 loss:  0.20356236398220062\n",
      "epoch:  53 loss:  0.20672015845775604\n",
      "epoch:  54 loss:  0.2022150605916977\n",
      "epoch:  55 loss:  0.20257779955863953\n",
      "epoch:  56 loss:  0.20376458764076233\n",
      "epoch:  57 loss:  0.20168638229370117\n",
      "epoch:  58 loss:  0.20051829516887665\n",
      "epoch:  59 loss:  0.20297928154468536\n",
      "epoch:  60 loss:  0.19942258298397064\n",
      "epoch:  61 loss:  0.19980232417583466\n",
      "epoch:  62 loss:  0.198652982711792\n",
      "epoch:  63 loss:  0.1982552409172058\n",
      "epoch:  64 loss:  0.2007521241903305\n",
      "epoch:  65 loss:  0.20180441439151764\n",
      "epoch:  66 loss:  0.19781090319156647\n",
      "epoch:  67 loss:  0.19678588211536407\n",
      "epoch:  68 loss:  0.19964908063411713\n",
      "epoch:  69 loss:  0.19736343622207642\n",
      "epoch:  70 loss:  0.1961665153503418\n",
      "epoch:  71 loss:  0.19671039283275604\n",
      "epoch:  72 loss:  0.1959608793258667\n",
      "epoch:  73 loss:  0.1963002234697342\n",
      "epoch:  74 loss:  0.19520051777362823\n",
      "epoch:  75 loss:  0.1953948587179184\n",
      "epoch:  76 loss:  0.19388407468795776\n",
      "epoch:  77 loss:  0.19388267397880554\n",
      "epoch:  78 loss:  0.19614766538143158\n",
      "epoch:  79 loss:  0.19351676106452942\n",
      "epoch:  80 loss:  0.19489333033561707\n",
      "epoch:  81 loss:  0.19259503483772278\n",
      "epoch:  82 loss:  0.19300739467144012\n",
      "epoch:  83 loss:  0.19165538251399994\n",
      "epoch:  84 loss:  0.1975150853395462\n",
      "epoch:  85 loss:  0.19308649003505707\n",
      "epoch:  86 loss:  0.1932307481765747\n",
      "epoch:  87 loss:  0.19102248549461365\n",
      "epoch:  88 loss:  0.1902049481868744\n",
      "epoch:  89 loss:  0.19086632132530212\n",
      "epoch:  90 loss:  0.19126303493976593\n",
      "epoch:  91 loss:  0.19162483513355255\n",
      "epoch:  92 loss:  0.19887924194335938\n",
      "epoch:  93 loss:  0.1894875019788742\n",
      "epoch:  94 loss:  0.191417396068573\n",
      "epoch:  95 loss:  0.19240480661392212\n",
      "epoch:  96 loss:  0.18841996788978577\n",
      "epoch:  97 loss:  0.1909097284078598\n",
      "epoch:  98 loss:  0.1936829835176468\n",
      "epoch:  99 loss:  0.18904533982276917\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for x,y in HR_dl:\n",
    "        y_gred = model(x)\n",
    "        loss = loss_fn(y_gred,y)\n",
    "        optim.zero_grad()\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "    with torch.no_grad():\n",
    "        print('epoch: ',epoch,'loss: ',loss_fn(model(X),Y).data.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c14f9d9-7325-47ac-aa2a-930b559cf120",
   "metadata": {},
   "source": [
    "# 使用 sklearn 柴分训练数据与测试数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "daa9f31f-6c6b-4c69-b7c1-847069e57c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c400f7d6-845e-4b3b-a7b5-3da5633bb5d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 默认是划分25%作为验证数据\n",
    "train_x,test_x,train_y,test_y = train_test_split(X,Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1739cab8-96bf-40a6-9fe4-7fb53fdca23a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([14999, 20]), torch.Size([14999, 1]))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape,Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a9f34f51-31d3-41e9-b81c-9701dec85755",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([11249, 20]),\n",
       " torch.Size([3750, 20]),\n",
       " torch.Size([11249, 1]),\n",
       " torch.Size([3750, 1]))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_x.shape,test_x.shape,train_y.shape,test_y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "602ebe74-92b4-419d-86ab-a0de6d6528f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([11249, 20])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_x.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "40b199b3-89bf-43e2-9657-f69f204d042a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.25001666777785186"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "3750 / 14999"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "58673c3d-63ae-46dd-9f6f-8e9941967401",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = TensorDataset(train_x,train_y)\n",
    "train_dl = DataLoader(train_ds,batch_size=batch,shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c05b1975-129a-4a1f-b243-d6dcb2293185",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = TensorDataset(test_ds_x,test_ds_y)\n",
    "test_dl = DataLoader(test_ds,batch_size=batch)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
